import torch
import torch.nn as nn


class MarginCalibrationModule(nn.Module):
    def __init__(self, fc):
        """
        Initializes the MarginCalibrationModule with trainable parameters omega and beta.

        Args:
        output_dim (int): The size of the output layer for which calibration is being applied.
        """
        super(MarginCalibrationModule, self).__init__()
        self.fc = fc
        self.weight = fc.weight
        self.nb_classes = fc.nb_classes
        # freeze
        self.weight.requires_grad = False
        for param in self.fc.parameters():
            param.requires_grad = False

        self.omega = nn.Parameter(torch.ones(1, self.nb_classes))
        self.beta = nn.Parameter(torch.zeros(1, self.nb_classes))

    def forward(self, feats):
        """
        Forward pass through the calibration module.

        Args:
            logits (torch.Tensor): The logits to be calibrated.

        Returns:
            torch.Tensor: Calibrated logits.
        """
        with torch.no_grad():
            w_norm = torch.norm(self.weight, dim=1, keepdim=False)
            logit_before = self.fc.forward(feats)["logits"]
        logit_after = self.omega * logit_before + self.beta * w_norm
        return {"logits": logit_after}
